# %%
import math
import torch
import torch.nn as nn
# %%
class MultiHeadSelfAttentionFormal(nn.Module):
    def __init__(self, hidden_dim: int, head_num: int, attention_dropout_rate: float=0.1) -> None:
        super().__init__()
        self.hidden_dim = hidden_dim
        self.head_num = head_num

        self.head_dim = hidden_dim // head_num

        self.q = nn.Linear(hidden_dim, hidden_dim) 
        self.k = nn.Linear(hidden_dim, hidden_dim) 
        self.v = nn.Linear(hidden_dim, hidden_dim)

        self.out = nn.Linear(hidden_dim, hidden_dim)

        self.attention_dropout = nn.Dropout(attention_dropout_rate)

    def forward(self, x, attention_mask=None):

        batch, seq_len, _ = x.size()

        Q = self.q(x)
        K = self.k(x)
        V = self.v(x)

        # (b, s, h) => (b, head_num, s, head_dim) # (h => head_num * head_dim)
        q_state = Q.view(batch, seq_len, self.head_num, self.head_dim).transpose(1, 2)
        k_state = K.view(batch, seq_len, self.head_num, self.head_dim).transpose(1, 2)
        v_state = V.view(batch, seq_len, self.head_num, self.head_dim).transpose(1, 2)

        # (b, head_num, s, s)
        attention_weight = q_state @ k_state.transpose(-1, 2) / math.sqrt(self.head_dim)

        if attention_mask is not None:
            attention_weight = attention_weight.masked_fill(
                attention_mask == 0,
                float('-inf')
            )
        print("attention_weight:", attention_weight)

        attention_weight = torch.softmax(attention_weight, dim=-1)
        attention_weight = self.attention_dropout(attention_weight)

        output_mid = torch.matmul(attention_weight, v_state)

        output_mid = output_mid.transpose(1, 2).contiguous()
        output_mid = output_mid.view(batch, seq_len, -1)
        out = self.out(output_mid)
        return out
# %%
